{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# N-Link Pendulum Control"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Taken from [PyDy](https://github.com/pydy/pydy/tree/master/examples/npendulum). This is not actually an example of symforce but rather a notable achievment with the underlying tools."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Several pieces of the puzzle have come together lately to really demonstrate the power of the scientific python software packages to handle complex dynamic and controls problems (i.e. IPython notebooks, matplotlib animations, python-control, and our software packages: sympy.physics.mechanics and PyDy).\n",
    "\n",
    "This [blog post by Wolfram](http://blog.wolfram.com/2011/03/01/stabilized-n-link-pendulum/) demonstrates Mathematica's ability to symbolically derive the equations of motion for the n-link pendulum and stabilize it with an LQR controller. This blog post inspired us to replicate the example with all free and open source software.\n",
    "\n",
    "In this example problem, we derive the equations of motion of an n-link pendulum on a laterally sliding cart and then develop a controller to stabilize it. Balancing a single inverted pendulum is a classic problem that is often a student's first experience with non-linear dynamics and control. The problem here is extended to a general n-link pendulum in which the equations of motion quickly get messy with greater than 2 links.\n",
    "\n",
    "The diagram below shows the general description of the problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![svg-pend](https://raw.githubusercontent.com/pydy/pydy/master/examples/npendulum/n-pendulum-with-cart.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Equations of Motion\n",
    "===================\n",
    "\n",
    "We'll start by generating the equations of motion for the system with SymPy **[mechanics](http://docs.sympy.org/dev/modules/physics/mechanics/index.html)**. The functionality that mechanics provides is much more in depth than Mathematica's functionality. In the Mathematica example, Lagrangian mechanics were implemented manually with Mathematica's symbolic functionality. **mechanics** provides an assortment of functions and classes to derive the equations of motion for arbitrarily complex (i.e. configuration constraints, nonholonomic motion constraints, etc) multibody systems in a very natural way. First we import the necessary functionality from SymPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sm\n",
    "import sympy.physics.mechanics as me"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can enable mathematical rendering of the resulting equations in the notebook with the following command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "me.init_vprinting()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now specify the number of links, $n$. I'll start with 5 since the Wolfram folks only showed four."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**mechanics** will need the generalized coordinates, generalized speeds, and the input force which are all time dependent variables and the bob masses, link lengths, and acceleration due to gravity which are all constants. Time, $t$, is also made available because we will need to differentiate with respect to time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q = me.dynamicsymbols(\"q:{}\".format(n + 1))  # Generalized coordinates\n",
    "u = me.dynamicsymbols(\"u:{}\".format(n + 1))  # Generalized speeds\n",
    "f = me.dynamicsymbols(\"f\")  # Force applied to the cart\n",
    "\n",
    "m = sm.symbols(\"m:{}\".format(n + 1))  # Mass of each bob\n",
    "l = sm.symbols(\"l:{}\".format(n))  # Length of each link\n",
    "g, t = sm.symbols(\"g t\")  # Gravity and time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create and inertial reference frame $I$ and define the point, $O$, as the origin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "I = me.ReferenceFrame(\"I\")  # Inertial reference frame\n",
    "O = me.Point(\"O\")  # Origin point\n",
    "O.set_vel(I, 0)  # Origin's velocity is zero"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Secondly, we define the define the first point of the pendulum as a particle which has mass. This point can only move laterally and represents the motion of the \"cart\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P0 = me.Point(\"P0\")  # Hinge point of top link\n",
    "P0.set_pos(O, q[0] * I.x)  # Set the position of P0\n",
    "P0.set_vel(I, u[0] * I.x)  # Set the velocity of P0\n",
    "Pa0 = me.Particle(\"Pa0\", P0, m[0])  # Define a particle at P0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can define the $n$ reference frames, particles, gravitational forces, and kinematical differential equations for each of the pendulum links. This is easily done with a loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames = [I]  # List to hold the n + 1 frames\n",
    "points = [P0]  # List to hold the n + 1 points\n",
    "particles = [Pa0]  # List to hold the n + 1 particles\n",
    "forces = [\n",
    "    (P0, f * I.x - m[0] * g * I.y)\n",
    "]  # List to hold the n + 1 applied forces, including the input force, f\n",
    "kindiffs = [q[0].diff(t) - u[0]]  # List to hold kinematic ODE's\n",
    "\n",
    "for i in range(n):\n",
    "    Bi = I.orientnew(\"B\" + str(i), \"Axis\", [q[i + 1], I.z])  # Create a new frame\n",
    "    Bi.set_ang_vel(I, u[i + 1] * I.z)  # Set angular velocity\n",
    "    frames.append(Bi)  # Add it to the frames list\n",
    "\n",
    "    Pi = points[-1].locatenew(\"P\" + str(i + 1), l[i] * Bi.x)  # Create a new point\n",
    "    Pi.v2pt_theory(points[-1], I, Bi)  # Set the velocity\n",
    "    points.append(Pi)  # Add it to the points list\n",
    "\n",
    "    Pai = me.Particle(\"Pa\" + str(i + 1), Pi, m[i + 1])  # Create a new particle\n",
    "    particles.append(Pai)  # Add it to the particles list\n",
    "\n",
    "    forces.append((Pi, -m[i + 1] * g * I.y))  # Set the force applied at the point\n",
    "\n",
    "    kindiffs.append(q[i + 1].diff(t) - u[i + 1])  # Define the kinematic ODE:  dq_i / dt - u_i = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With all of the necessary point velocities and particle masses defined, the `KanesMethod` class can be used to derive the equations of motion of the system automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kane = me.KanesMethod(I, q_ind=q, u_ind=u, kd_eqs=kindiffs)  # Initialize the object\n",
    "fr, frstar = kane.kanes_equations(particles, forces)  # Generate EoM's fr + frstar = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The equations of motion are quite long as can been seen below. This is the general nature of most non-simple mutlibody problems. That is why a SymPy is so useful; no more mistakes in algebra, differentiation, or copying hand written equations. Note that `trigsimp` can take quite a while to complete for extremely large expressions. Below we print $\\tilde{M}$ and $\\tilde{f}$ from $\\tilde{M}\\dot{u}=\\tilde{f}$ to show the size of the expressions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.trigsimp(kane.mass_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\tilde{M}$ is a function of the constant parameters and the configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "me.find_dynamicsymbols(kane.mass_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.trigsimp(kane.forcing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\tilde{f}$ is a function of the constant parameters, configuration, speeds, and the applied force."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "me.find_dynamicsymbols(kane.forcing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulation\n",
    "==========\n",
    "\n",
    "Now that the symbolic equations of motion are available we can simulate the pendulum's motion. We will need some more SymPy functionality and several NumPy functions, and most importantly the integration function from SciPy, `odeint`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy.linalg import solve\n",
    "from scipy.integrate import odeint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, define some numeric values for all of the constant parameters in the problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arm_length = 1.0 / n  # The maximum length of the pendulum is 1 meter\n",
    "bob_mass = 0.01 / n  # The maximum mass of the bobs is 10 grams\n",
    "parameters = [g, m[0]]  # Parameter definitions starting with gravity and the first bob\n",
    "parameter_vals = [9.81, 0.01 / n]  # Numerical values for the first two\n",
    "for i in range(n):  # Then each mass and length\n",
    "    parameters += [l[i], m[i + 1]]\n",
    "    parameter_vals += [arm_length, bob_mass]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mathematica has a really nice `NDSolve` function for quickly integrating their symbolic differential equations. We make use of SymPy's lambdify function to do something similar, i.e. to create functions that will evaluate the \"full\" mass matrix, $M$, and \"full\" forcing vector, $f$ from $M\\dot{x} = f(x, r, t)$ as a NumPy function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dynamic = q + u  # Make a list of the states\n",
    "dynamic.append(f)  # Add the input force\n",
    "\n",
    "M_func = sm.lambdify(\n",
    "    dynamic + parameters, kane.mass_matrix_full\n",
    ")  # Create a callable function to evaluate the mass matrix\n",
    "f_func = sm.lambdify(\n",
    "    dynamic + parameters, kane.forcing_full\n",
    ")  # Create a callable function to evaluate the forcing vector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To integrate the ODE's we need to define a function that returns the derivatives of the states given the current state and time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def right_hand_side(x, t, args):\n",
    "    \"\"\"\n",
    "    Returns the derivatives of the states.\n",
    "\n",
    "    Args:\n",
    "        x: ndarray, shape(2 * (n + 1))\n",
    "            The current state vector.\n",
    "        t: float\n",
    "            The current time.\n",
    "        args: ndarray\n",
    "            The constants.\n",
    "\n",
    "    Returns:\n",
    "        dx: ndarray, shape(2 * (n + 1))\n",
    "            The derivative of the state.\n",
    "    \"\"\"\n",
    "    r = 0.0  # The input force is always zero\n",
    "    arguments = np.hstack((x, r, args))  # States, input, and parameters\n",
    "    dx = np.array(solve(M_func(*arguments), f_func(*arguments))).T[0]  # Solving for the derivatives\n",
    "\n",
    "    return dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the right hand side function, the initial conditions are set such that the pendulum is in the vertical equilibrium and a slight initial rate is set for each speed to ensure the pendulum falls. The equations can then be integrated with SciPy's `odeint` function given a time series."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = np.hstack(\n",
    "    (0.0, np.pi / 2 * np.ones(len(q) - 1), 1e-3 * np.ones(len(u)))  # q0  # q1...qn+1\n",
    ")  # u0...un+1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.linspace(0.0, 10.0, num=500)  # Time vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = odeint(right_hand_side, x0, t, args=(parameter_vals,))  # Numerical integration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting\n",
    "========\n",
    "\n",
    "The results of the simulation can be plotted with matplotlib. First, load the plotting functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.pylabtools import figsize\n",
    "\n",
    "figsize(8.0, 6.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The coordinate trajectories are plotted below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = plt.plot(t, x[:, : x.shape[1] // 2])\n",
    "lab = plt.xlabel(\"Time [sec]\")\n",
    "leg = plt.legend(dynamic[: x.shape[1] // 2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the generalized speed trajectories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = plt.plot(t, x[:, x.shape[1] // 2 :])\n",
    "lab = plt.xlabel(\"Time [sec]\")\n",
    "leg = plt.legend(dynamic[x.shape[1] // 2 :])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Animation\n",
    "=========\n",
    "\n",
    "matplotlib now includes very nice animation functions for animating matplotlib plots. First we import the necessary functions for creating the animation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import animation\n",
    "from matplotlib.patches import Rectangle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following function was modeled from Jake Vanderplas's [post on matplotlib animations](http://jakevdp.github.com/blog/2012/08/18/matplotlib-animation-tutorial/). The default animation writer is used (typically ffmpeg), you can change it by adding `writer` argument to `anim.save` call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def animate_pendulum(t, states, length, filename=None):\n",
    "    \"\"\"\n",
    "    Animates the n-pendulum and optionally saves it to file.\n",
    "\n",
    "    Args:\n",
    "        t: ndarray, shape(m)\n",
    "            Time array.\n",
    "        states: ndarray, shape(m,p)\n",
    "            State time history.\n",
    "        length: float\n",
    "            The length of the pendulum links.\n",
    "        filename: string or None, optional\n",
    "            If true a movie file will be saved of the animation. This may take some time.\n",
    "\n",
    "    Returns:\n",
    "        fig: matplotlib.Figure\n",
    "            The figure.\n",
    "        anim: matplotlib.FuncAnimation\n",
    "            The animation.\n",
    "    \"\"\"\n",
    "    # the number of pendulum bobs\n",
    "    numpoints = states.shape[1] // 2\n",
    "\n",
    "    # first set up the figure, the axis, and the plot elements we want to animate\n",
    "    fig = plt.figure()\n",
    "\n",
    "    # some dimesions\n",
    "    cart_width = 0.4\n",
    "    cart_height = 0.2\n",
    "\n",
    "    # set the limits based on the motion\n",
    "    xmin = np.around(states[:, 0].min() - cart_width / 2.0, 1)\n",
    "    xmax = np.around(states[:, 0].max() + cart_width / 2.0, 1)\n",
    "\n",
    "    # create the axes\n",
    "    ax = plt.axes(xlim=(xmin, xmax), ylim=(-1.1, 1.1), aspect=\"equal\")\n",
    "\n",
    "    # display the current time\n",
    "    time_text = ax.text(0.04, 0.9, \"\", transform=ax.transAxes)\n",
    "\n",
    "    # create a rectangular cart\n",
    "    rect = Rectangle(\n",
    "        [states[0, 0] - cart_width / 2.0, -cart_height / 2],\n",
    "        cart_width,\n",
    "        cart_height,\n",
    "        fill=True,\n",
    "        color=\"red\",\n",
    "        ec=\"black\",\n",
    "    )\n",
    "    ax.add_patch(rect)\n",
    "\n",
    "    # blank line for the pendulum\n",
    "    (line,) = ax.plot([], [], lw=2, marker=\"o\", markersize=6)\n",
    "\n",
    "    # initialization function: plot the background of each frame\n",
    "    def init():\n",
    "        time_text.set_text(\"\")\n",
    "        rect.set_xy((0.0, 0.0))\n",
    "        line.set_data([], [])\n",
    "        return time_text, rect, line\n",
    "\n",
    "    # animation function: update the objects\n",
    "    def animate(i):\n",
    "        time_text.set_text(\"time = {:2.2f}\".format(t[i]))\n",
    "        rect.set_xy((states[i, 0] - cart_width / 2.0, -cart_height / 2))\n",
    "        x = np.hstack((states[i, 0], np.zeros((numpoints - 1))))\n",
    "        y = np.zeros((numpoints))\n",
    "        for j in np.arange(1, numpoints):\n",
    "            x[j] = x[j - 1] + length * np.cos(states[i, j])\n",
    "            y[j] = y[j - 1] + length * np.sin(states[i, j])\n",
    "        line.set_data(x, y)\n",
    "        return time_text, rect, line\n",
    "\n",
    "    # call the animator function\n",
    "    anim = animation.FuncAnimation(\n",
    "        fig,\n",
    "        animate,\n",
    "        frames=len(t),\n",
    "        init_func=init,\n",
    "        interval=t[-1] / len(t) * 1000,\n",
    "        blit=True,\n",
    "        repeat=False,\n",
    "    )\n",
    "\n",
    "    # save the animation if a filename is given\n",
    "    if filename is not None:\n",
    "        anim.save(filename, fps=30, codec=\"libx264\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create the animation of the pendulum. This animation will show the open loop dynamics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "animate_pendulum(t, x, arm_length, filename=\"open-loop.mp4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "\n",
    "html = \"\"\"\n",
    "<video width=\"640\" height=\"480\" controls>\n",
    "  <source src=\"open-loop.mp4\" type=\"video/mp4\">\n",
    "Your browser does not support the video tag, check out the YouTube version instead: http://youtu.be/Nj3_npq7MZI.\n",
    "</video>\n",
    "\"\"\"\n",
    "HTML(html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Controller Design\n",
    "=================\n",
    "\n",
    "The n-link pendulum can be balanced such that all of the links are inverted above the cart by applying the correct lateral force to the cart. We can design a full state feedback controller based from a linear model of the pendulum about its upright equilibrium point. We'll start by specifying the equilibrium point and parameters in dictionaries. We make sure to use SymPy types in the equilibrium point to ensure proper cancelations in the linearization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "equilibrium_point = [sm.S(0)] + [sm.pi / 2] * (len(q) - 1) + [sm.S(0)] * len(u)\n",
    "equilibrium_dict = dict(zip(q + u, equilibrium_point))\n",
    "equilibrium_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `KanesMethod` class has method that linearizes the forcing vector about generic state and input perturbation vectors. The equilibrium point and numerical constants can then be substituted in to give the linear system in this form: $M\\dot{x}=F_Ax+F_Br$. The state and input matrices, $A$ and $B$, can then be computed by left side multiplication by the inverse of the mass matrix: $A=M^{-1}F_A$ and $B=M^{-1}F_B$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M, F_A, F_B, r = kane.linearize(new_method=True, op_point=equilibrium_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.simplify(M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.simplify(F_A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm.simplify(F_B).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the numerical $A$ and $B$ matrices can be formed. First substitute numerical parameter values into $M$, $F_A$, and $F_B$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameter_dict = dict(zip(parameters, parameter_vals))\n",
    "parameter_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M_num = sm.matrix2numpy(M.subs(parameter_dict), dtype=float)\n",
    "F_A_num = sm.matrix2numpy(F_A.subs(parameter_dict), dtype=float)\n",
    "F_B_num = sm.matrix2numpy(F_B.subs(parameter_dict), dtype=float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = np.linalg.solve(M_num, F_A_num)\n",
    "B = np.linalg.solve(M_num, F_B_num)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also convert `equilibrium_point` to a numeric array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "equilibrium_point = np.asarray([x.evalf() for x in equilibrium_point], dtype=float)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have a linear system, the SciPy package can be used to design an optimal controller for the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.linalg import solve_continuous_are"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we can check to see if the system is, in fact, controllable. The rank of the controllability matrix must be equal to the number of rows in $A$, but the `matrix_rank` algorithm is numerically ill conditioned and for certain values of $n$ this will fail, as seen below for $n=5$. Nevertheless, the system is controllable, no matter the number of links."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def controllable(a, b):\n",
    "    \"\"\"\n",
    "    Returns true if the system is controllable and false if not.\n",
    "\n",
    "    Args:\n",
    "        a: array_like, shape(n,n)\n",
    "            The state matrix.\n",
    "        b: array_like, shape(n,r)\n",
    "            The input matrix.\n",
    "\n",
    "    Returns:\n",
    "        controllable: boolean\n",
    "    \"\"\"\n",
    "    a = np.matrix(a)\n",
    "    b = np.matrix(b)\n",
    "    n = a.shape[0]\n",
    "    controllability_matrix = []\n",
    "    for i in range(n):\n",
    "        controllability_matrix.append(a**i * b)\n",
    "    controllability_matrix = np.hstack(controllability_matrix)\n",
    "\n",
    "    return np.linalg.matrix_rank(controllability_matrix, tol=1e-6) == n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "controllable(A, B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So now we can compute the optimal gains with a linear quadratic regulator. I chose identity matrices for the weightings for simplicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = np.eye(A.shape[0])\n",
    "R = np.eye(B.shape[1])\n",
    "S = solve_continuous_are(A, B, Q, R)\n",
    "K = np.dot(np.dot(np.linalg.inv(R), B.T), S)\n",
    "K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The gains can now be used to define the required input during simulation to stabilize the system. The input $r$ is simply the gain vector multiplied by the error in the state vector from the equilibrium point, $r(t)=K(x_{eq} - x(t))$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def right_hand_side(x, t, args):\n",
    "    \"\"\"\n",
    "    Returns the derivatives of the states.\n",
    "\n",
    "    Args:\n",
    "        x: ndarray, shape(2 * (n + 1))\n",
    "            The current state vector.\n",
    "        t: float\n",
    "            The current time.\n",
    "        args: ndarray\n",
    "            The constants.\n",
    "\n",
    "    Returns:\n",
    "        dx: ndarray, shape(2 * (n + 1))\n",
    "            The derivative of the state.\n",
    "    \"\"\"\n",
    "    r = np.dot(K, equilibrium_point - x)  # The controller\n",
    "    arguments = np.hstack((x, r, args))  # States, input, and parameters\n",
    "    dx = np.array(solve(M_func(*arguments), f_func(*arguments))).T[0]  # Solving for the derivatives\n",
    "\n",
    "    return dx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can simulate and animate the system to see if the controller works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = np.hstack((0, np.pi / 2 * np.ones(len(q) - 1), 1 * np.ones(len(u))))\n",
    "t = np.linspace(0.0, 10.0, num=500)\n",
    "x = odeint(right_hand_side, x0, t, args=(parameter_vals,))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The plots show that we seem to have a stable system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = plt.plot(t, x[:, : x.shape[1] // 2])\n",
    "lab = plt.xlabel(\"Time [sec]\")\n",
    "leg = plt.legend(dynamic[: x.shape[1] // 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = plt.plot(t, x[:, x.shape[1] // 2 :])\n",
    "lab = plt.xlabel(\"Time [sec]\")\n",
    "leg = plt.legend(dynamic[x.shape[1] // 2 :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "animate_pendulum(t, x, arm_length, filename=\"closed-loop.mp4\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "\n",
    "html = \"\"\"\n",
    "<video width=\"640\" height=\"480\" controls>\n",
    "  <source src=\"closed-loop.mp4\" type=\"video/mp4\">\n",
    "Your browser does not support the video tag, check out the YouTube version instead: http://youtu.be/SpgBHqW9om0\n",
    "</video>\n",
    "\"\"\"\n",
    "HTML(html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The video clearly shows that the controller can balance all $n$ of the pendulum links. The weightings in the lqr design can be tweaked to give different performance if needed.\n",
    "\n",
    "This example shows that the free and open source scientific Python tools for dynamics are easily comparable in ability and quality to a commercial package such as Mathematica. \n",
    "\n",
    "The IPython notebook for this example can be downloaded from https://github.com/pydy/pydy/tree/master/examples/npendulum. You can try out different $n$ values. I've gotten the equations of motion to compute for an open loop simulation of 10 links. My computer ran out of memory when I tried to compute for $n=50$. The controller weightings and initial conditions will probably have to be adjusted for better performance for $n>5$, but it should work."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
